

#ifndef KDTREE_H
#define KDTREE_H

// STL includes
#include <string>
#include <cstring>
#include <iostream>
using std::string;

#include <vector>
using std::vector;

#include <algorithm>

#include "Geometry.h"
#include "HitData.h"

class Ray;
class KdTree;
struct HitData;


enum SplitEventTypes 
{ 
	SPLIT_START, 
	SPLIT_END 
};

struct KdSplitEvent
{
	float m_position;
	enum SplitEventTypes m_type;
	
	// Distance comparison for intersection segment sorting
	friend bool operator<(const KdSplitEvent& left, const KdSplitEvent& right);	
};

inline bool operator<(const KdSplitEvent& left, const KdSplitEvent& right)
{
	return left.m_position < right.m_position;
}

struct KdSplittingPlane
{
	int m_axis;       // normal axis to spltting plane: x=0, y=1, z=2
	float m_position; // plane position on this axis
	
	enum Type
	{
		IN_BACK_OF, 
		IN_FRONT_OF, 
		SPANNING
	};
};

struct KdSpanningStats
{
	int m_nFront;
	int m_nBack;
	int m_nSpanning;
	
};

class KdNode
{
public:
	
	KdNode() : m_front(NULL), m_back(NULL)
	{
		
	}
	
	// Each node stores the indices of the triangles in KdTree.m_geometry which it contains
	vector<int> m_triangleIndices;
	
	KdSplittingPlane m_splittingPlane;
	Aabb m_aabb;
	
	struct KdNode* m_front;
	struct KdNode* m_back;
	
	void build(const vector<int>& triangleIndicesIn, KdTree& parentTree, int treeDepth);
	
	void draw(int currentDepth, int maxDepth);
	
	void intersect(const Aabb& aabb, vector<int>& triangleIndicesOut);

private:
	
	void _check();
	void _updateLeafStats( KdTree& parentTree, int treeDepth );
	bool _findBestSplittingPlane( KdTree& parentTree );
	void _getSubAaabbs( Aabb& frontAabb, Aabb& backAabb, const Aabb&inAabb, const KdSplittingPlane& splitting );
	float _computeSplittingCost( Aabb& frontAabb, Aabb& backAabb, Aabb& nodeAabb, KdSpanningStats& stats );
	float _SAH(const KdSplittingPlane& candidateSplittingPlane, int NL, int NR, int NS);
	KdSplittingPlane::Type _classifyTriangle(const Aabb& triangleAabb);
	void _getFrontAndBackLists(vector<int>& frontList, vector<int>& backList, KdTree& parentTree);
};



class KdTree
{	
	
public:
			
	/* On construction, builds the KD tree from the provided world space geometry. */
	KdTree(const Geometry& geometry);

	virtual ~KdTree();		

	
	/// QUERY
	
	// Find which triangles of the geometry overlap the given AABB
	void intersect(const Aabb& aabb, vector<int>& triangleIndicesOut);
	
	// Intersect a Ray with the geometry, returning the first intersection (or an empty intersection if nothing was hit).
	HitData intersect(Ray& ray) const;
	
	// Draw a UI representation of the kd-tree
	void draw(int maxDepth);

	void reportStats();

	// Build parameters
	int m_maxDepth;
	float m_maxCostRatio;
	
	// Build stats
	int m_cumulativeLeafDepth;
	int m_numLeaves;
	
	// Structure used during runtime traversal
	struct StackElmt
	{
		class KdNode* m_node;
		
		float m_t;  // the entry/exit signed distance
		Vec4 m_pb;  // the coordinates of entry/exit point
		int m_prev; // the pointer to the previous stack item
	};
	
	friend class KdNode;
	
private:
	
	struct KdNode* m_root;
	
	// Geometry indexed into by the KD-tree nodes.
	struct Geometry m_geometry;

	HitData m_noHit;
			
};



#endif //KDTREE_H




